from collections import Counter
from datetime import datetime, date

from rest_framework import viewsets, mixins
from rest_framework.decorators import action
from rest_framework.views import Response, status
from rest_framework.permissions import IsAuthenticated

from apps.common.account.models import Account
from apps.pollution.models import CompanySelfAssessment


class IndexViewSet(mixins.ListModelMixin, viewsets.GenericViewSet):
    permission_classes = [IsAuthenticated]

    def list(self, request, *args, **kwargs):
        """ 获取首页数据 """
        need_company_qs = Account.objects.filter(is_company=True, is_analysis=True).values_list("username", flat=True)
        need_company_cnt = len(need_company_qs)
        qs = CompanySelfAssessment.objects.filter(company_name__in=need_company_qs).values_list('check_status', flat=True)
        counter = Counter(qs)
        no_view_cnt = counter[0]
        no_problem_cnt = counter[1]
        had_problem_cnt = counter[2]
        all_cnt = len(qs)
        no_write_cnt = need_company_cnt - all_cnt

        res_data = {
            'need_company_cnt': need_company_cnt, 'no_view_cnt': no_view_cnt, 'no_problem_cnt': no_problem_cnt,
            'had_problem_cnt': had_problem_cnt, 'all_cnt': all_cnt, 'no_write_cnt': no_write_cnt
        }
        return Response(res_data, status=status.HTTP_200_OK)
